import torch
from torch.utils.data import Dataset
from torchvision.transforms import transforms
import cv2
import os
from copy import deepcopy

class MyDataset(Dataset):
    def __init__(self, img_rootPath, mode='train', scale=0.9):
        super().__init__()
        data_dir = []

        for dir in os.listdir(img_rootPath):
            for img in os.listdir(os.path.join(img_rootPath, dir)):
                img_path = os.path.join(img_rootPath, dir, img)
                if os.path.splitext(img_path)[-1] == '.jpg':
                    data_dir.append(img_path)

        # 数据集划分
        if mode == 'train':
            # data_dir = data_dir[:int(len(data_dir)*0.01)]
            data_dir = data_dir[:int(len(data_dir)*scale)]
        elif mode == 'test':
            # data_dir = data_dir[:int(len(data_dir)*0.01)]
            data_dir = data_dir[int(len(data_dir)*(scale)):]

        self.img_path_arr = data_dir

        self.size = 128

        self.transform = transforms.Compose([
            transforms.ToTensor()
        ])
    
    def __getitem__(self, item):
        # 读取图片
        img_path = self.img_path_arr[item]
        img = cv2.imread(img_path)
        img = cv2.resize(img, [self.size, self.size])

        # 转Tensor并归一化
        img = self.transform(img)
        return img

    def __len__(self):
        return len(self.img_path_arr)


if __name__ == '__main__':
    dataset = MyDataset(r'D:\VOCtrainval_11-May-2012\JPEGImages')
    print(dataset[5][0])